#include "CustomTexture.h"
#include "../../Utils/LogUtil.h"
#include "../../Core/CoreGlobals.h"
#include "../RenderState.h"
#include "../../Game/Process/RenderProcess.h"

CustomTexture::BuilderDetails::BuilderDetails()
{
	isMultiple = false;

	flags = 0;
	imageType = VK_IMAGE_TYPE_2D;
	format = VK_FORMAT_B8G8R8A8_UNORM;
	extent.width = GRenderState->GetWidth();
	extent.height = GRenderState->GetHeight();
	extent.depth = 1;
	mipLevels = 1;
	arrayLayers = 1;
	samples = VK_SAMPLE_COUNT_1_BIT;
	tiling = VK_IMAGE_TILING_OPTIMAL;
	usage = VK_IMAGE_USAGE_SAMPLED_BIT;
	sharingMode = VK_SHARING_MODE_EXCLUSIVE;
	imageLayout = VK_IMAGE_LAYOUT_UNDEFINED;
	mMemoryProperty = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;

	viewType = VK_IMAGE_VIEW_TYPE_2D;
	components = { VK_COMPONENT_SWIZZLE_R , VK_COMPONENT_SWIZZLE_G , VK_COMPONENT_SWIZZLE_B , VK_COMPONENT_SWIZZLE_A };
	subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
	subresourceRange.baseMipLevel = 0;
	subresourceRange.levelCount = VK_REMAINING_MIP_LEVELS;
	subresourceRange.baseArrayLayer = 0;
	subresourceRange.layerCount = VK_REMAINING_ARRAY_LAYERS;

	data = nullptr;
}

CustomTexture::Builder::Builder() noexcept = default;
CustomTexture::Builder::~Builder() noexcept = default;
CustomTexture::Builder::Builder(CustomTexture::Builder const& rhs) noexcept = default;
CustomTexture::Builder::Builder(CustomTexture::Builder && rhs) noexcept = default;
CustomTexture::Builder& CustomTexture::Builder::operator=(CustomTexture::Builder const& rhs) noexcept = default;
CustomTexture::Builder& CustomTexture::Builder::operator=(CustomTexture::Builder && rhs) noexcept = default;


CustomTexture::Builder& CustomTexture::Builder::SetMultiple(bool isMultiple) noexcept
{
	mImpl->isMultiple = isMultiple;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetFlags(VkImageCreateFlags flags) noexcept
{
	mImpl->flags = flags;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetImageType(VkImageType imageType) noexcept
{
	mImpl->imageType = imageType;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetFormat(VkFormat format) noexcept
{
	mImpl->format = format;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetExtent(VkExtent3D extent) noexcept
{
	mImpl->extent = extent;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetWidth(uint32 width) noexcept
{
	mImpl->extent.width = width;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetHeight(uint32 height) noexcept
{
	mImpl->extent.height = height;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetDepth(uint32 depth) noexcept
{
	mImpl->extent.depth = depth;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetMipLevels(uint32 mipLevels) noexcept
{
	mImpl->mipLevels = mipLevels;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetArrayLayers(uint32 arrayLayers) noexcept
{
	mImpl->arrayLayers = arrayLayers;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetSamples(VkSampleCountFlagBits samples) noexcept
{
	assert(samples <= GVulkanRHI->GetMsaaSamples());
	mImpl->samples = samples;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetTiling(VkImageTiling tiling) noexcept
{
	mImpl->tiling = tiling;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetUsage(VkImageUsageFlags usage) noexcept
{
	mImpl->usage = usage;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetSharingMode(VkSharingMode sharingMode) noexcept
{
	mImpl->sharingMode = sharingMode;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetImageLayout(VkImageLayout imageLayout) noexcept
{
	mImpl->imageLayout = imageLayout;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetMemoryProperty(VkMemoryPropertyFlags mMemoryProperty) noexcept
{
	mImpl->mMemoryProperty = mMemoryProperty;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::AddQueueFamily(VkQueueFlagBits queueFlag) noexcept
{
	mImpl->queueFamilys.push_back(queueFlag);
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetViewType(VkImageViewType viewType) noexcept
{
	mImpl->viewType = viewType;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetComponents(VkComponentMapping components) noexcept
{
	mImpl->components = components;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetAspectMask(VkImageAspectFlags aspectMask) noexcept
{
	mImpl->subresourceRange.aspectMask = aspectMask;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetBaseMipLevel(uint32 baseMipLevel) noexcept
{
	mImpl->subresourceRange.baseMipLevel = baseMipLevel;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetLevelCount(uint32 levelCount) noexcept
{
	mImpl->subresourceRange.levelCount = levelCount;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetBaseArrayLayer(uint32 baseArrayLayer) noexcept
{
	mImpl->subresourceRange.baseArrayLayer = baseArrayLayer;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetLayerCount(uint32 layerCount) noexcept
{
	mImpl->subresourceRange.layerCount = layerCount;
	return *this;
}

CustomTexture::Builder& CustomTexture::Builder::SetData(void* data) noexcept
{
	mImpl->data = data;
	return *this;
}

RefCountPtr<CustomTexture> CustomTexture::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<CustomTexture>(NewObject<CustomTexture>(*mImpl));
}

RefCountPtr<CustomTexture> CustomTexture::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	CustomTexture* customTexture = NewObject<CustomTexture>(*mImpl);
	process->AddTexture(name, customTexture);
	return RefCountPtr<CustomTexture>(customTexture);
}

void CustomTexture::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveTexture(mName, this);
		}
		DeleteObject<CustomTexture>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveTexture(mName, this);
			}
			DeleteObject<CustomTexture>(this);
			});
	}
}

CustomTexture::CustomTexture(BuilderDetails & builderDetails)
{
	mIsMultiple = builderDetails.isMultiple;
	mIsRenderTarget = false;

	mFlags = builderDetails.flags;
	mImageType = builderDetails.imageType;
	mFormat = builderDetails.format;
	mExtent = builderDetails.extent;
	mMipLevels = builderDetails.mipLevels;
	mArrayLayers = builderDetails.arrayLayers;
	mSamples = builderDetails.samples;
	mTiling = builderDetails.tiling;
	mUsage = builderDetails.usage;
	mSharingMode = builderDetails.sharingMode;
	mQueueFamilys.assign(builderDetails.queueFamilys.begin(), builderDetails.queueFamilys.end());
	if (mQueueFamilys.size() > 1)
	{
		// Check if the queueFamilys is same
		bool isDifferent = false;
		const uint32 standQueueFamily = GVulkanRHI->GetQueueFamily(mQueueFamilys[0]);
		for (uint32 i = 1; i < mQueueFamilys.size(); ++i)
		{
			if (GVulkanRHI->GetQueueFamily(mQueueFamilys[i]) != standQueueFamily)
				isDifferent = true;
		}

		if (isDifferent)
			mSharingMode = VK_SHARING_MODE_CONCURRENT; // multi queue
		else
			mSharingMode = VK_SHARING_MODE_EXCLUSIVE;
	}
	else
	{
		assert(mSharingMode == VK_SHARING_MODE_EXCLUSIVE);
	}
	mImageLayouts.resize(GetFrameCount(), builderDetails.imageLayout);
	mMemoryProperty = builderDetails.mMemoryProperty;

	mViewType = builderDetails.viewType;
	mComponents = builderDetails.components;
	mSubresourceRange = builderDetails.subresourceRange;
	if (mSubresourceRange.levelCount == VK_REMAINING_MIP_LEVELS)
	{
		mSubresourceRange.levelCount = mMipLevels;
	}
	if (mSubresourceRange.layerCount == VK_REMAINING_ARRAY_LAYERS)
	{
		mSubresourceRange.layerCount = mArrayLayers;
	}

	mImages.resize(GetFrameCount(), VK_NULL_HANDLE);
	mImageMemorys.resize(GetFrameCount(), VK_NULL_HANDLE);
	mImageViews.resize(GetFrameCount(), VK_NULL_HANDLE);

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		GVulkanRHI->CreateImage(mImages[i], mImageMemorys[i], mFlags, mExtent, mUsage, mFormat, mMipLevels, mMemoryProperty, mSamples, mImageType, mImageLayouts[i], mArrayLayers, mTiling, mSharingMode, mQueueFamilys);
	}


	// if the data is not null, create the buffer to copy data to image
	if (builderDetails.data != nullptr)
	{
		// If you want to copy data, need to be transfer dst bit
		assert(mUsage & VK_IMAGE_USAGE_TRANSFER_DST_BIT);

		VkBuffer stagingBuffer;
		VkDeviceMemory stagingMemory;

		VkDeviceSize dataSize = static_cast<VkDeviceSize>(mExtent.width * mExtent.height * mExtent.depth * GVkFormatSize[mFormat]);

		GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, dataSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

		void* data;
		VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, dataSize, 0, &data));
		memcpy(data, builderDetails.data, dataSize);
		vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

		for (uint32 i = 0; i < GetFrameCount(); ++i)
		{
			// TransitionImageLayout for copy data
			GVulkanRHI->TransferLayout(commandBuffer, mImages[i], mImageLayouts[i], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
			mImageLayouts[i] = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;

			GVulkanRHI->CopyBufferToImage(commandBuffer, stagingBuffer, mImages[i], mExtent, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, mSubresourceRange.layerCount });

			if (mMipLevels > 1)
			{
				// If you want to generate mipmap, need to copy data from image, use VK_IMAGE_USAGE_TRANSFER_SRC_BIT
				assert(mUsage & VK_IMAGE_USAGE_TRANSFER_SRC_BIT != 0);
				GVulkanRHI->GenerateMipmaps2D(commandBuffer, mImages[i], mExtent.width, mExtent.height, mFormat, mMipLevels, mArrayLayers);
			}
			else
			{
				GVulkanRHI->TransferLayout(commandBuffer, mImages[i], mImageLayouts[i], VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, mSubresourceRange, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT);
			}
			mImageLayouts[i] = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
		}

		GRenderState->FlushSingleCommand(commandBuffer);

		vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
		vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
	}

	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		GVulkanRHI->CreateImageView(mImageViews[i], mImages[i], mFormat, mViewType, mSubresourceRange, mComponents);
	}
}

